/*
 * @Author: Zhou Zijian
 * @Date: 2023-02-28 23:32:29
 * @Last Modified by: Zhou Zijian
 * @Last Modified time: 2023-03-06 02:31:14
 */

#ifndef GPU_X_H
#define GPU_X_H

#include <CL/cl.h>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>

using clGetPlatformIDsPtrFunc = decltype(&clGetPlatformIDs);
using clGetDeviceIDsPtrFunc = decltype(&clGetDeviceIDs);
using clCreateContextPtrFunc = decltype(&clCreateContext);
using clCreateCommandQueuePtrFunc = decltype(&clCreateCommandQueue);
using clCreateProgramWithSourcePtrFunc = decltype(&clCreateProgramWithSource);
using clBuildProgramPtrFunc = decltype(&clBuildProgram);
using clCreateKernelPtrFunc = decltype(&clCreateKernel);
using clReleaseKernelPtrFunc = decltype(&clReleaseKernel);
using clReleaseProgramPtrFunc = decltype(&clReleaseProgram);
using clReleaseCommandQueuePtrFunc = decltype(&clReleaseCommandQueue);
using clReleaseContextPtrFunc = decltype(&clReleaseContext);
using clSetKernelArgPtrFunc = decltype(&clSetKernelArg);
using clEnqueueNDRangeKernelPtrFunc = decltype(&clEnqueueNDRangeKernel);
using clFinishPtrFunc = decltype(&clFinish);
using clCreateBufferPtrFunc = decltype(&clCreateBuffer);
using clEnqueueReadBufferPtrFunc = decltype(&clEnqueueReadBuffer);
using clReleaseMemObjectPtrFunc = decltype(&clReleaseMemObject);


class GpuX {
public:
    enum ErrorCode {
        OPEN_LIB_FAILED = -24,
        LOAD_FUNC_FAILED,
        CLOSE_LIB_FAILED,
        CONFIG_FAILED,
        MALLOC_FAILED,
        NO_PLATFORM,
        NO_GPU,
        GET_PLATFORM_ID_FAILED,
        GET_DEVICE_ID_FAILED,
        CREATE_CONTEXT_FAILED,
        CREATE_COMMAND_QUEUE_FAILED,
        FINISH_FAILED,
        CREATE_BUFFER_FAILED_NULLPTR,
        CREATE_BUFFER_FAILED_SIZE_ZERO,
        CREATE_BUFFER_FAILED,
        FLUSH_BUFFER_FAILED,
        NO_BIN_OR_CL,
        OPEN_FILE_FAILED,
        CREATE_PROGRAM_FAILED,
        BUILD_PROGRAM_FAILED,
        CREATE_KERNEL_FAILED,
        SET_KERNEL_ARG_FAILED,
        FUNCTION_NOT_EXIST,
        RUN_KERNEL_FAILED,
        SUCCESS
    };

private:
    class OpenclLibLoader {
    public:
        int Open(const char *libName);
        int Close();

    public:
        clGetPlatformIDsPtrFunc GetPlatformIDs = nullptr;
        clGetDeviceIDsPtrFunc GetDeviceIDs = nullptr;
        clCreateContextPtrFunc CreateContext = nullptr;
        clCreateCommandQueuePtrFunc CreateCommandQueue = nullptr;
        clCreateProgramWithSourcePtrFunc CreateProgramWithSource = nullptr;
        clBuildProgramPtrFunc BuildProgram = nullptr;
        clCreateKernelPtrFunc CreateKernel = nullptr;
        clReleaseKernelPtrFunc ReleaseKernel = nullptr;
        clReleaseProgramPtrFunc ReleaseProgram = nullptr;
        clReleaseCommandQueuePtrFunc ReleaseCommandQueue = nullptr;
        clReleaseContextPtrFunc ReleaseContext = nullptr;
        clSetKernelArgPtrFunc SetKernelArg = nullptr;
        clEnqueueNDRangeKernelPtrFunc EnqueueNDRangeKernel = nullptr;
        clFinishPtrFunc Finish = nullptr;
        clCreateBufferPtrFunc CreateBuffer = nullptr;
        clEnqueueReadBufferPtrFunc EnqueueReadBuffer = nullptr;
        clReleaseMemObjectPtrFunc ReleaseMemObject = nullptr;
    
    private:
        void *m_handle = nullptr;
    };

public:
    class Buffer {
    public:
        Buffer(OpenclLibLoader *cl, const cl_context &context, const cl_command_queue &commandQueue)
            : m_cl(cl), m_context(context), m_commandQueue(commandQueue)
        {}
        ~Buffer();
        int Init(void *ptr, size_t size);
        cl_mem GetClMem();
        void *GetPointer();
        size_t GetSize();
        int Flush();

    private:
        OpenclLibLoader *m_cl;
        cl_context m_context;
        cl_command_queue m_commandQueue;
        void *m_ptr = nullptr;
        size_t m_size = 0;
        cl_mem m_mem;
    };

public:
    static GpuX &GetInstance()
    {
        static GpuX instance;
        return instance;
    }
    int Create(const std::string &config);
    void Destroy();
    int GetClError();
    int Finish();
    Buffer *CreateBuffer(void *ptr, size_t size);
    void DestroyBuffer(Buffer *buffer);
    template <typename T, typename... Ts>
    int Run(const std::string &function, bool async, T arg, Ts... args)
    {
        auto findKernel = m_functions.find(function);
        auto findGlobalWorkSize = m_globalWorkSize.find(function);
        auto findLocalWorkSize = m_localWorkSize.find(function);
        if (findKernel == m_functions.end() || findGlobalWorkSize == m_globalWorkSize.end() ||
            findLocalWorkSize == m_localWorkSize.end()) {
            return GpuX::ErrorCode::FUNCTION_NOT_EXIST;
        }
        cl_kernel kernel = findKernel->second;
        int ret = SetKernelArg(kernel, 0, arg, args...);
        if (ret != GpuX::ErrorCode::SUCCESS) {
            return ret;
        }
        auto globalWorkSize = findGlobalWorkSize->second;
        auto localWorkSize = findLocalWorkSize->second;
        cl_uint workDim = globalWorkSize.size();
        m_err = m_cl.EnqueueNDRangeKernel(m_commandQueue, kernel, workDim, nullptr, globalWorkSize.data(),
            localWorkSize.data(), 0, nullptr, nullptr);
        if (m_err != CL_SUCCESS) {
            return GpuX::ErrorCode::RUN_KERNEL_FAILED;
        }
        if (!async) {
            ret = Finish();
            if (ret != GpuX::ErrorCode::SUCCESS) {
                return ret;
            }
        }
        return GpuX::ErrorCode::SUCCESS;
    }

private:
    GpuX() = default;
    ~GpuX();
    int GetDeviceId();
    int GetContext();
    int GetCommandQueue();
    int BuildProgram(const std::string &program);
    int SetFunctions(const std::string &name, const std::string &kernel, const std::vector<size_t> &globalWorkSize,
        const std::vector<size_t> &localWorkSize);
    template <typename T, typename... Ts>
    int SetKernelArg(const cl_kernel &kernel, int index, T arg, Ts... args)
    {
        m_err = m_cl.SetKernelArg(kernel, index, sizeof(arg), &arg);
        if (m_err != CL_SUCCESS) {
            return GpuX::ErrorCode::SET_KERNEL_ARG_FAILED;
        }
        if constexpr (sizeof...(args) > 0) {
            return SetKernelArg(kernel, index + 1, args...);
        }
        return GpuX::ErrorCode::SUCCESS;
    }

public:
    GpuX(const GpuX &) = delete;
    GpuX &operator=(const GpuX) = delete;

private:
    OpenclLibLoader m_cl;
    cl_int m_err = CL_SUCCESS;
    cl_device_id m_deviceId{0};
    cl_context m_context{0};
    cl_command_queue m_commandQueue{0};
    cl_program m_program{0};
    std::set<cl_kernel> m_kernels;
    std::unordered_map<std::string, cl_kernel> m_functions;
    std::unordered_map<std::string, std::vector<size_t>> m_globalWorkSize;
    std::unordered_map<std::string, std::vector<size_t>> m_localWorkSize;
};

#endif  // GPU_X_H
